# chapter 5 种子生成器规模化自动生成代码
# 加入对生成种子的质量筛选，选择效果最好的种子生成器。
import ast
import sys
import os

from openai import OpenAI
key = "" #换成你的key
url = "" # 换成你的url
model_name = "" # 换成你的模型名，建议使用qwen3-32b以上的模型
def run_llm(prompt_str):
    result = ""
    client = OpenAI(api_key=key, base_url=url)
    response = client.chat.completions.create(
        model=model_name,
                    messages=[
                        {"role": "user", "content": prompt_str},
                    ],
                    stream=True,
                    # 把timeout值设置高一些
                    timeout=600
                )
    for chunk in response:
        content = chunk.choices[0].delta.content
        if content is not None:
            result += content
            print_llm(content)
    return result


wrapper_code = """


target_floder = "./j40/corpus"
run(target_floder, 40)
"""

# 大语言模型回显部分
global print_count
print_count = 0
def print_llm(text):
    global print_count
    symbols = ['|','|','|','|','|', '/', '/', '/', '/', '/', '-','-','-','-','-', '\\','\\','\\','\\','\\']
    print_count += 1
    print_count = print_count % len(symbols)
    sys.stdout.write('\r大语言模型工作中 ' + symbols[print_count])
    sys.stdout.flush()
    
    
def extract_python_code(prompt_str):
    import re
    python_code = re.findall(r"```python(.*?)```", prompt_str, re.DOTALL)[0]
    return python_code
# step1 将生成部分包装成函数
def gen_seed_generator_code(feature_result):
    # 生成python代码

    code_gen_prompt_str = f"""根据下面的jxl格式要求，写一个Python脚本，用来生成满足下面要求的jxl文件,要求要使用pillow_jxl库。
    要求：1.输出的文件夹要用户指定。
    2. 输出的函数接口名为run,里面要有输出文件夹、输出文件数量参数两个参数。
    3. 整个程序只有一个run函数，所有生成逻辑都要放到run函数中。
    4. 除了pillow_jxl库外，只能使用struct，numpy,os,sys库，严禁使用其他库。
    5. 不需要加入日志或者其他与生成jxl文件无关的功能。
    6. 不要加入任何调用run函数的示例代码。

    {feature_result}
    生成的代码是python代码：
    ```python
    //代码放到这里
    ```
    \\no_think
    """
    result = run_llm(code_gen_prompt_str)
    return result
def generate_featrue():
    # 生成特性
    feature_prompt_str = """根据你对jxl图像知识，列出jxl图像在文件格式上的10种算法特性。
    \\no_think
    """
    feature_result = run_llm(feature_prompt_str)
    return feature_result


# 获取覆盖率
def get_cov():
    run_target_cmd = "./j40/test ./j40/corpus -runs=2 -detect_leaks=0"
    os.system(run_target_cmd)
    os.system("llvm-profdata merge -sparse *.profraw -o default.profdata")
    get_cov_cmd = "llvm-cov report ./j40/test -instr-profile=./default.profdata > cov_info"
    os.system(get_cov_cmd)
    with open("cov_info", "r") as f:
        cov_info = f.read()
    cov_line = cov_info.split("TOTAL")[1]
    cov_list = cov_line.split(" ")
    cov = cov_list[-1].split("%")[0]
    os.system("rm cov_info")
    return float(cov)


# 先生成特性
feature_result = generate_featrue()

# step 2 设立死循环，不断尝试生成代码，直到成功为止
target_floder = "./j40/corpus"
rm_cmd =f"rm {target_floder}/*"

# 设置全局变量
best_python_code = ""
success_gen_count = 0
MAX_GEN_COUNT = 10
MAX_COV = 0
while True:
    print("\n代码生成中...")
    # 删除之前的种子文件
    for filename in os.listdir(target_floder):
        file_path = os.path.join(target_floder, filename)
        os.remove(file_path)
    rough_code = gen_seed_generator_code(feature_result)
    print("\n")
    python_code = extract_python_code(rough_code)
    # 如果无法解析到python代码，则重新生成
    if not python_code: continue
    run_code = python_code + wrapper_code
    # 解析目标语法是否有问题
    try:
        ast.parse(run_code)
    except:
        print("\nInvalid Python Code\n")
        continue
    # 尝试运行目标代码，如果成功，则检查是否生成了种子
    try:
        exec(run_code)
        file_list = os.listdir(target_floder)
        if len(file_list) > 0:
            print("\n种子生成成功！\n")
            success_gen_count += 1
            cov = get_cov()
            print(f"当前总覆盖率: {cov}。 目前最大的覆盖率: {MAX_COV}")
            if cov > MAX_COV:
                best_python_code = run_code
                MAX_COV = cov
            if success_gen_count > MAX_GEN_COUNT:
                break
        else:
            print("\n生成的代码没有产生种子\n")
    except:
        print("\n生成代码运行报错\n")
        continue
    
print(f"最终结果，最大覆盖率为: {MAX_COV}")
print("最好的代码:")
print(best_python_code)
with open("best_python_code.py","w") as f:
    f.write(best_python_code)